#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <time.h>
#include <gmp.h>
#include <stdbool.h>

#define MAX_MSG 100
#define MAX_K 16384
#define MIN_K 512
#define ERROR_INT -1
#define PRINT_BASE 62
#define KEY_LEN 2048
#define LEN_BIT2OCTET(x) x/8

typedef unsigned char UCHAR;
typedef unsigned int SIZE_T;
typedef unsigned long SIZE_BIT_T;

int generateKeys(const SIZE_BIT_T k, mpz_t *kn, mpz_t *ke, mpz_t *kd);
int RSA_PKCS1_v1_5_Encrypt(const UCHAR *M, const mpz_t n, const mpz_t e, const SIZE_T k, UCHAR **C);
int RSA_PKCS1_v1_5_Decrypt(const UCHAR *C, const mpz_t n, const mpz_t d, const SIZE_T k, UCHAR **M);
void EME_PKCS1_v1_5_Encoding(UCHAR *EM, const UCHAR *M, const SIZE_T k, const SIZE_T mLen);
void EME_PKCS1_v1_5_Decoding(UCHAR *C, const UCHAR *EM, const SIZE_T emLen);
void OS2IP(mpz_t *x, const UCHAR *X, const SIZE_T xLen);
void I2OSP(UCHAR *X, mpz_t x, const SIZE_T xLen);
void RSAEP(mpz_t *c, const mpz_t m, const mpz_t n, const mpz_t e);
void RSADP(mpz_t *m, const mpz_t c, const mpz_t n, const mpz_t d);
void randStr(UCHAR *s, const SIZE_T len);
void printmpzlen(char *msg, const SIZE_T alen, const mpz_t a);
void printmpz(char *msg, const mpz_t a);
int mystrcat(char *dest, const char *src, const SIZE_T ptr, const SIZE_T src_len);
int main()
{
    UCHAR M[MAX_MSG];
    UCHAR *C = NULL;
    UCHAR *plain = NULL;
    int k = LEN_BIT2OCTET(KEY_LEN);
	mpz_t n, e, d;
    mpz_inits(n, e, d, NULL);

	generateKeys(KEY_LEN, &n, &e, &d);
    printmpz("n:", n);
    printmpz("e:", e);
    printmpz("d:", d);
    printf("\n");

    scanf("%s", M);
	printf("\nplaintext: %s\n\n", M);

	RSA_PKCS1_v1_5_Encrypt(M, n, e, k, &C);
	printf("ciphertext: %s\n\n", C);
    
	RSA_PKCS1_v1_5_Decrypt(C, n, d, k, &plain);
	printf("decypted text: %s\n", plain);

    free(C);
    free(plain);
    mpz_clears(n, e, d, NULL);
}

int generateKeys(const SIZE_BIT_T k, mpz_t *kn, mpz_t *ke, mpz_t *kd){
    //检查k是否合法
    if(k < MIN_K|| k > MAX_K){
        return ERROR_INT;
    }
    mpz_t p, q, n, d, e, p_tmp, q_tmp, p_m1, q_m1, phi_n, gcd;
    mpz_inits(p, q, n, d, e, p_tmp, q_tmp, p_m1, q_m1, phi_n, gcd, NULL);
    
	//设置e为常数值65537
    mpz_set_ui(e, 65537);
	//设置随机数种子
    gmp_randstate_t state;
    gmp_randinit_default(state);
    clock_t tm = time(NULL);
    gmp_randseed_ui(state, tm);
	//设定p、q比特数
    mp_bitcnt_t p_bitcnt = (k + 1) >> 1;
    mp_bitcnt_t q_bitcnt = k - p_bitcnt;
    //生成p_bitcnt比特随机整数p_tmp
    mpz_rrandomb(p_tmp, state, p_bitcnt);
    //寻找合适的p、q、
    while(1) {
        //生成n的其中一个质因子p
        //p为大于p_tmp的下一个质数
        mpz_nextprime(p, p_tmp);
        //生成k bit的大整数n
        do {
            //生成n的另一个质因子q
            gmp_randseed_ui(state, ++tm);
            mpz_rrandomb(q_tmp, state, q_bitcnt);
            mpz_nextprime(q, q_tmp);
            if(mpz_cmp(p, q) < 0) {
                mpz_swap(p, q);
            }
            //计算n
            mpz_mul(n, p, q);
        } while(mpz_sizeinbase(n, 2) < k);
		//计算φ(n)=(p-1)(q-1)
        mpz_sub_ui(p_m1, p, 1);
        mpz_sub_ui(q_m1, q, 1);
        mpz_mul(phi_n, p_m1, q_m1);
        //计算gcd(φ(n), e)
        mpz_gcd(gcd, phi_n, e);
        //若gcd(φ(n), e) != 1，则重新寻找p、q
        if (mpz_cmp_ui(gcd, 1) != 0) continue;
		
        //计算d为e的模φ(n)逆元
        mpz_invert(d, e, phi_n);
        //返回n，e，d
        mpz_set(*kn, n);
        mpz_set(*ke, e);
        mpz_set(*kd, d);
        break;
    }

    mpz_clears(p, q, n, d, e, p_tmp, q_tmp, p_m1, q_m1, phi_n, gcd, NULL);
    return 0;
}

int RSA_PKCS1_v1_5_Encrypt(const UCHAR *M, const mpz_t n, const mpz_t e, const SIZE_T k, UCHAR **C){
    int mLen = strlen(M);
    if(mLen > k - 11){
        printf("message too long\n");
        return ERROR_INT;
    }

    UCHAR *EM = (UCHAR*)malloc(k+1);
    *C = (UCHAR*)malloc(k+1);
    mpz_t m, c;
    memset(EM, '\0', k+1);
    memset(*C, '\0', k+1);
    mpz_inits(m, c, NULL);

    EME_PKCS1_v1_5_Encoding(EM, M, k, mLen);
    OS2IP(&m, EM, k);
    RSAEP(&c, m, n, e);
    I2OSP(*C, c, k);

    mpz_clears(m, c, NULL);
    free(EM);
    return 0;
}

int RSA_PKCS1_v1_5_Decrypt(const UCHAR *C, const mpz_t n, const mpz_t d, const SIZE_T k, UCHAR **M){
    UCHAR *EM = (UCHAR*)malloc(k+1);
    *M = (UCHAR*)malloc(k+1);
    mpz_t c, m;
    mpz_inits(c, m, NULL);
    memset(EM, '\0', k+1);
    memset(*M, '\0', k+1);

    OS2IP(&c, C, k);
    RSADP(&m, c, n, d);
    I2OSP(EM, m, k); 
    EME_PKCS1_v1_5_Decoding(*M, EM, k);

    free(EM);
    mpz_clears(c, m, NULL);
    return 0;
}

void EME_PKCS1_v1_5_Encoding(UCHAR *EM, const UCHAR *M, const SIZE_T k, const SIZE_T mLen){
    UCHAR prefix[] = {0, 2};
    UCHAR suffix[] = {0};
    UCHAR *ps = (char*)malloc(k - mLen - 3);
    memset(ps, 0, k-mLen-3);
    randStr(ps, k-mLen-3);

    EM[0] = prefix[0];
    EM[1] = prefix[1];
    int i = 2;
    i = mystrcat(EM, ps, i, k-mLen-3);
    i = mystrcat(EM, suffix, i, 1);
    i = mystrcat(EM, M, i, mLen);
    free(ps);
}

void EME_PKCS1_v1_5_Decoding(UCHAR *M, const UCHAR *EM, const SIZE_T emLen){
    int i = 2, j = 0;
    bool is_msg = false;

    while(i < emLen){
        if(is_msg){
            M[j] = EM[i];
            j++;
        }
        if(EM[i] == 0){
            is_msg = true;
        }
        i++;
    }
}

void OS2IP(mpz_t *x, const UCHAR *X, const SIZE_T xLen){
    int i=1;
    mpz_t tmp, pow, base, xi;
    mpz_inits(tmp, pow, base, xi, NULL);
    mpz_set_ui(*x, 0);
    mpz_set_ui(base, 256);
    while(i <= xLen) {
        mpz_pow_ui(pow, base, (xLen- i));
        mpz_mul_ui(xi, pow, X[i-1]);
        mpz_add(tmp, tmp, xi);
        i++;
    }
    mpz_set(*x, tmp);
    mpz_clears(tmp, pow, base, xi, NULL);
}

void I2OSP(UCHAR *X, mpz_t x, const SIZE_T xLen) {
    mpz_t uplim, base, pow, xi, remain;
    mpz_inits(uplim, base, pow, xi, remain, NULL);
    mpz_set(remain, x);
    mpz_set_ui(base, 256);
    mpz_pow_ui(uplim, base, xLen);

    if (mpz_cmp(x, uplim) >= 0){
        printf("integer too large\n");
        mpz_clears(uplim, base, pow, xi, NULL);
        return;
    }

    int i = 1;
    while(i <= xLen) {
        mpz_pow_ui(pow, base, (xLen - i));
        mpz_tdiv_q(xi, remain, pow);
        mpz_mod(remain, x, pow);
        X[i-1] = (UCHAR)mpz_get_ui(xi);
        i++;
    }

    mpz_clears(uplim, base, pow, xi, remain, NULL);
}

void RSAEP(mpz_t *c, const mpz_t m, const mpz_t n, const mpz_t e){
    mpz_t up_lim, down_lim;
    mpz_inits(up_lim, down_lim, NULL);
    mpz_set_ui(down_lim, 0);
    mpz_sub_ui(up_lim, n, 1);

    if(mpz_cmp(m, down_lim) < 0 || mpz_cmp(m, up_lim) > 0){
        mpz_clears(up_lim, down_lim, NULL);
        printf("message representive out of range\n");
        return;
    }
    mpz_powm(*c, m, e, n);
    mpz_clears(up_lim, down_lim, NULL);
} 

void RSADP(mpz_t *m, const mpz_t c, const mpz_t n, const mpz_t d){
    mpz_t up_lim, down_lim;
    mpz_inits(up_lim, down_lim, NULL);
    mpz_set_ui(down_lim, 0);
    mpz_sub_ui(up_lim, n, 1);

    if(mpz_cmp(c, down_lim) < 0 || mpz_cmp(c, up_lim) > 0){
        mpz_clears(up_lim, down_lim, NULL);
        printf("message representive out of range\n");
        return;
    }
    mpz_powm(*m, c, d, n);
    mpz_clears(up_lim, down_lim, NULL);
}

void randStr(UCHAR *s, const SIZE_T len){
    int i = 0;
    srand((unsigned) time(NULL));
    while(i < len) {
        s[i] = rand() % 93 + 33;
        i++;
    }
}

void printmpzlen(char *msg, const SIZE_T alen, const mpz_t a) {
    char *as = mpz_get_str(NULL, PRINT_BASE, a);
    printf("%s\n", msg);
    printf("len:%d\n", alen);
    printf("%s\n", as);
}

void printmpz(char *msg, const mpz_t a){
    char *as = mpz_get_str(NULL, PRINT_BASE, a);
    printf("%s ", msg);
    printf("%s\n", as);
}

int mystrcat(char *dest, const char *src, const SIZE_T ptr, const SIZE_T src_len){
    int p = ptr, i = 0;
    while(i < src_len){
        dest[p] = src[i];
        p++, i++;
    }
    return p;
}